"""
找出 1~40 之间最优的 K 值，训练 KNN 手写数字分类器，
并将最佳模型持久化到 best_knn.pickle
"""
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score

# ---------- 1. 加载数字数据集 ----------
digits = load_digits()
X, y = digits.data, digits.target      # 1 797 样本，64 维像素特征

# ---------- 2. 划分训练集 / 测试集 ----------
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42, stratify=y)

# ---------- 3. 初始化变量 ----------
best_k, best_score, best_model = 1, 0.0, None
k_range = range(1, 41)
scores = []

# ---------- 4. 网格搜索最优 K ----------
for k in tqdm(k_range, desc="Tuning K"):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    pred = knn.predict(X_test)
    acc = accuracy_score(y_test, pred)
    scores.append(acc)
    if acc > best_score:
        best_score, best_k, best_model = acc, k, knn

# ---------- 5. 保存最佳模型 ----------
with open("best_knn.pickle", "wb") as f:
    pickle.dump(best_model, f)

# ---------- 6. 打印结果 ----------
print(f"Best K = {best_k}  |  Test Accuracy = {best_score:.4f}")

# ---------- 7. 绘制 K-acc 曲线（可选） ----------
plt.plot(k_range, scores, marker="o")
plt.title("KNN K-value vs Accuracy")
plt.xlabel("K")
plt.ylabel("Accuracy")
plt.grid(True)
plt.savefig("k_acc_curve.png")
plt.show()